#include "SSAOData.h"

//To get MSAA data - textures which are bound as RTVs need to match backbuffer
#include "GraphicsManager.h"

//Include effects we use
#include "SSAOMapFX.h"
#include "SSAOBlurFX.h"
#include "SSAOBlurCSFX.h"

//Also include shader manager to get access to them
#include "ShaderManager.h"

//Include the camera
#include "VirtualCamera.h"

SSAOData::SSAOData(ID3D11Device* device, unsigned height, unsigned width, 
	ID3D11Buffer* quadVB, ID3D11Buffer* quadIB)
{
	//Ensure that device, height and width are valid values. 
	assert(height > 0);
	assert(width > 0);
	assert(device != 0);
	//Assert quads are legit. 
	assert(quadVB);
	assert(quadIB);

	//Init views to NULL. 
	viewNormalDepthRTV = 0;
	viewNormalDepthSRV = 0;
	ssaoMapRTV = 0;
	ssaoMapSRV = 0;
	ssaoMapRandomVectorsSRV = 0;
	ssaoMapGenVB = 0;
	ssaoMapGenIB = 0;
	ssaoMapRTV1 = 0;
	ssaoMapSRV1 = 0;

	//UAV for CS
	ssaoMapUAV = 0;
	ssaoMapUAV1 = 0;

	//Store width/height
	rtHeight = height;
	rtWidth = width;

	//Get effects from the shader manager. 
	ShaderManager* sharedShaderManager = ShaderManager::GetInstance();
	ssaoMapFX = sharedShaderManager->ReturnSSAOMapFX();
	ssaoBlurFX = sharedShaderManager->ReturnSSAOBlurFX();
	ssaoBlurCSFX = sharedShaderManager->ReturnSSAOBlurCSFX();

	//Creates views. 
	CreateViews(device);

	//Create quads. 
	CreateSSAOQuad(device);

	//Generate offset vectors and random vector texture for pass 2
	CreateOffsetVectors();
	CreateRandomVectorTexture(device);

	//Store pointers to quad buffers 
	this->quadVB = quadVB;
	this->quadIB = quadIB;
}

SSAOData::~SSAOData()
{
	//Destroy views
	//
	//Phase 1
	ReleaseCOM(viewNormalDepthRTV);
	ReleaseCOM(viewNormalDepthSRV);

	//Phase 2
	ReleaseCOM(ssaoMapRTV);
	ReleaseCOM(ssaoMapSRV);
	ReleaseCOM(ssaoMapRandomVectorsSRV);

	//Phase 3
	ReleaseCOM(ssaoMapRTV1);
	ReleaseCOM(ssaoMapSRV1);

	//Release quads
	ReleaseCOM(ssaoMapGenVB);
	ReleaseCOM(ssaoMapGenIB);

	//Release UAVS
	ReleaseCOM(ssaoMapUAV);
	ReleaseCOM(ssaoMapUAV1);
}

void SSAOData::GenerateSSAOMap(ID3D11DeviceContext* context, VirtualCamera* camera)
{
	//Store fov/farz. 
	fovy = camera->GetFovY();
	farz = camera->GetFarZ();

	//Rebuild frustum corners. 
	BuildFrustumCorners(fovy, farz);

	//Bind ssaoMapRTV as render target. Disable depth/stencil view
	//as its not required. This will also disable depth tests.
	ID3D11RenderTargetView* rtv[1] = {ssaoMapRTV};
	context->OMSetRenderTargets(1, rtv, NULL);
	context->ClearRenderTargetView(ssaoMapRTV, reinterpret_cast<const float*>(&Colors::Black));
	context->RSSetViewports(1, &ssaoMapViewport);
	
	//Matrix that transforms from NDC space to texture space.
	static const XMMATRIX tex(
		0.5f, 0.0f, 0.0f, 0.0f,
		0.0f, -0.5f, 0.0f, 0.0f,
		0.0f, 0.0f, 1.0f, 0.0f,
		0.5f, 0.5f, 0.0f, 1.0f);

	XMMATRIX viewToTex = XMMatrixMultiply(camera->GetProjection(), tex);

	//Execute draw command. This will render our SSAO factors
	//in to the render target.(ssaoMapRTV) 
	context->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
	context->IASetInputLayout(ssaoMapFX->ReturnInputLayout());

	UINT stride = sizeof(SSAOMapFX::SSAOMapVertex);
	UINT offset = 0;

	context->IASetVertexBuffers(0, 1, &ssaoMapGenVB, &stride, &offset);
	context->IASetIndexBuffer(ssaoMapGenIB, DXGI_FORMAT_R32_UINT, 0);

	//Set data
	ssaoMapFX->SetCBufferPerFrame(viewToTex, offsetVectors, frustumCorners);
	ssaoMapFX->SetShaderResources(viewNormalDepthSRV, ssaoMapRandomVectorsSRV);

	//Draw
	for (unsigned p = 0; p < ssaoMapFX->ReturnActiveTechniqueNumberOfPasses(); p++)
	{
		ssaoMapFX->ApplyActiveTechniqueForPass(p, context);
		context->DrawIndexed(6, 0, 0);
	}

	//Unbding SRVs
	ID3D11ShaderResourceView* vNULL[] = {0,0};
	context->PSSetShaderResources(0, 2, vNULL);
}

void SSAOData::BlurSSAOMap(ID3D11DeviceContext* context, unsigned blurCount)
{
	for (unsigned i = 0; i < blurCount; i++)
	{
		Blur(context, ssaoMapSRV, ssaoMapRTV1, true);
		Blur(context, ssaoMapSRV1, ssaoMapRTV, false);
	}
}

void SSAOData::BlurSSAOMapCS(ID3D11DeviceContext* context, unsigned blurCount)
{
	//Unbind RTV
	context->OMSetRenderTargets(0, 0, NULL);

	for (unsigned i = 0; i < blurCount; i++)
	{
		BlurCS(context, ssaoMapSRV, ssaoMapUAV1, true);
		BlurCS(context, ssaoMapSRV1, ssaoMapUAV, false);
	}
}

//Carries out the blur
void SSAOData::Blur(ID3D11DeviceContext* context,
	ID3D11ShaderResourceView* input, ID3D11RenderTargetView* output, bool horizontal)
{
	//Set render target to the output - again, no depth buffer to disable
	//depth test/writes.
	ID3D11RenderTargetView* rtv[1] = {output};
	context->OMSetRenderTargets(1, rtv, 0);
	context->ClearRenderTargetView(output, reinterpret_cast<const float*>(&Colors::Black));
	context->RSSetViewports(1, &ssaoMapViewport);

	//Set active technique based on the bool passed in.
	//
	//We have a differnt technique for each blur pass (horizontal/vertical) 
	if (horizontal)
		ssaoBlurFX->SetActiveTechniqueByName("SSAOBlurHorizontal");
	else
		ssaoBlurFX->SetActiveTechniqueByName("SSAOBlurVertical");

	//Setup IA stage.
	context->IASetPrimitiveTopology(D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST);
	context->IASetInputLayout(ssaoBlurFX->ReturnInputLayout());

	UINT stride = sizeof(SSAOBlurFX::SSAOBlurVertex);
	UINT offset = 0;

	context->IASetVertexBuffers(0, 1, &quadVB, &stride, &offset);
	context->IASetIndexBuffer(quadIB, DXGI_FORMAT_R32_UINT, 0);

	//Set effect input. 
	ssaoBlurFX->SetShaderResources(input, viewNormalDepthSRV);
	//Set texel data
	ssaoBlurFX->SetCBufferTexelData(
		1.0f / (float)ssaoMapViewport.Width,
		1.0f / (float)ssaoMapViewport.Height);

	//Draw
	for (unsigned p = 0; p < ssaoBlurFX->ReturnActiveTechniqueNumberOfPasses(); p++)
	{
		ssaoBlurFX->ApplyActiveTechniqueForPass(p, context);
		context->DrawIndexed(6, 0, 0);
	}

	//Unbind input. 
	ssaoBlurFX->SetShaderResources(0, 0);
	ssaoBlurFX->ApplyActiveTechniqueForPass(0, context);
}

void SSAOData::BlurCS(ID3D11DeviceContext* context,
		ID3D11ShaderResourceView* input, ID3D11UnorderedAccessView* output, bool horizontal)
{
	//Work out how many thread groups to dispatch. 
	UINT dispatchX = (UINT)ceilf( ((float)rtWidth / 2.0f) / 32.0f);
	UINT dispatchY = (UINT)ceilf( ((float)rtHeight / 2.0f) / 32.0f);

	//Set resources.
	ssaoBlurCSFX->SetTextures(input, output);

	for (unsigned p = 0; p < ssaoBlurCSFX->ReturnActiveTechniqueNumberOfPasses(); p++)
	{
		//Dispatch
		ssaoBlurCSFX->ApplyActiveTechniqueForPass(p, context);
		if (horizontal)
			context->Dispatch(dispatchX, dispatchY, 1);
		else
			context->Dispatch(dispatchX, dispatchY, 1);
	}

	//Unbind resources.
	ID3D11ShaderResourceView* nullSRV[1] = {0};
	ID3D11UnorderedAccessView* nullUAV[1] = {0};

	context->CSSetShaderResources(0, 1, nullSRV);
	context->CSSetUnorderedAccessViews(0, 1, nullUAV, 0);
}

void SSAOData::Rebuild(ID3D11Device* device, unsigned newHeight, unsigned newWidth)
{
	//Ensuree valid device. 
	assert(device != 0);

	//Store new sizes
	rtHeight = newHeight;
	rtWidth = newWidth;

	//Destroy old data and recreate. 
	CreateViews(device);

	//Create viewport for pass two - 1/2 dimentions of the RT
	ssaoMapViewport.TopLeftX = 0.0f;
	ssaoMapViewport.TopLeftY = 0.0f;
	ssaoMapViewport.MinDepth = 0.0f;
	ssaoMapViewport.MaxDepth = 1.0f;
	ssaoMapViewport.Width  = (float)newWidth / 2.0f;
	ssaoMapViewport.Height = (float)newHeight / 2.0f;

	//Generate offset vectors and random vector texture for pass 2
	//Destroy old random vectors texture
	ReleaseCOM(ssaoMapRandomVectorsSRV);

	//Recreate
	CreateOffsetVectors();
	CreateRandomVectorTexture(device);
}

void SSAOData::CreateViews(ID3D11Device* device)
{
	//Release old buffers/views
	ReleaseCOM(viewNormalDepthRTV);
	ReleaseCOM(viewNormalDepthSRV);
	//Release SSAOMap RTV/SRVs
	ReleaseCOM(ssaoMapRTV);
	ReleaseCOM(ssaoMapSRV);
	//Release second SSAOMap srv/rtv
	ReleaseCOM(ssaoMapRTV1);
	ReleaseCOM(ssaoMapSRV1);
	//Release UAVS
	ReleaseCOM(ssaoMapUAV);
	ReleaseCOM(ssaoMapUAV1);

	//Create views for stage 1. 
	CreateNormalDepthViews(device);
	//Crreate views fro stage 2
	CreateSSAOMapTexture(device);
	//Create second views for stage 3. 
	CreateSecondSSAOMapTexture(device);
}

void SSAOData::CreateNormalDepthViews(ID3D11Device* device)
{
	//For MSAA data
	GraphicsManager* sharedGraphicsManager = GraphicsManager::GetInstance();

	//Create a texture2D object for our RTV and SRV - should match
	//the size of the client width/height.
	ID3D11Texture2D* tex;

	D3D11_TEXTURE2D_DESC ndDesc;
	ndDesc.ArraySize = 1;
	ndDesc.BindFlags = D3D11_BIND_RENDER_TARGET | D3D11_BIND_SHADER_RESOURCE;
	ndDesc.CPUAccessFlags = NULL;
	ndDesc.Format = DXGI_FORMAT_R16G16B16A16_FLOAT;
	ndDesc.Height = rtHeight;
	ndDesc.Width = rtWidth;
	ndDesc.MipLevels = 1;
	ndDesc.MiscFlags = 0;

	//Check MSAA
	if (sharedGraphicsManager->IsMSAAEnabled())
	{
		ndDesc.SampleDesc.Count = 4;
		ndDesc.SampleDesc.Quality = sharedGraphicsManager->GetMSAAQuality();
	}
	else
	{
		ndDesc.SampleDesc.Count = 1;
		ndDesc.SampleDesc.Quality = 0;
	}

	ndDesc.Usage = D3D11_USAGE_DEFAULT;

	//Create the texture.
	HR(device->CreateTexture2D(&ndDesc, 0, &tex));

	//Create RTV and SRV to this texture.
	HR(device->CreateRenderTargetView(tex, 0, &viewNormalDepthRTV));
	HR(device->CreateShaderResourceView(tex, 0, &viewNormalDepthSRV));

	//Release texture as the RTV and SRV increment the
	//reference count.
	ReleaseCOM(tex);
}

void SSAOData::CreateSSAOMapTexture(ID3D11Device* device)
{
	//Create a texture pointer
	ID3D11Texture2D* texture;

	//Fillout Tex2D description for the first
	//SSAOMap Texture
	D3D11_TEXTURE2D_DESC texDesc;
	texDesc.ArraySize = 1;
	texDesc.BindFlags = D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_RENDER_TARGET | D3D11_BIND_UNORDERED_ACCESS;
	texDesc.CPUAccessFlags = 0;
	texDesc.Format = DXGI_FORMAT_R16_FLOAT;
	texDesc.Height = rtHeight/2;
	texDesc.Width = rtWidth/2;
	texDesc.MipLevels = 1;
	texDesc.MiscFlags = 0;
	texDesc.SampleDesc.Count = 1;
	texDesc.SampleDesc.Quality = 0;

	texDesc.Usage = D3D11_USAGE_DEFAULT;

	//Create texture2D and the SRV/RTV
	HR(device->CreateTexture2D(&texDesc, 0, &texture));
	HR(device->CreateRenderTargetView(texture, 0, &ssaoMapRTV));
	HR(device->CreateShaderResourceView(texture, 0, &ssaoMapSRV));
	HR(device->CreateUnorderedAccessView(texture, 0, &ssaoMapUAV));

	//Done with the texture, release
	ReleaseCOM(texture);
}

void SSAOData::CreateSecondSSAOMapTexture(ID3D11Device* device)
{
	//Create a texture pointer
	ID3D11Texture2D* texture;

	//Fillout Tex2D description for the first
	//SSAOMap Texture
	D3D11_TEXTURE2D_DESC texDesc;
	texDesc.ArraySize = 1;
	texDesc.BindFlags = D3D11_BIND_SHADER_RESOURCE | D3D11_BIND_RENDER_TARGET | D3D11_BIND_UNORDERED_ACCESS;
	texDesc.CPUAccessFlags = 0;
	texDesc.Format = DXGI_FORMAT_R16_FLOAT;
	texDesc.Height = rtHeight/2;
	texDesc.Width = rtWidth/2;
	texDesc.MipLevels = 1;
	texDesc.MiscFlags = 0;
	texDesc.SampleDesc.Count = 1;
	texDesc.SampleDesc.Quality = 0;

	texDesc.Usage = D3D11_USAGE_DEFAULT;

	//Create texture2D and the SRV/RTV
	HR(device->CreateTexture2D(&texDesc, 0, &texture));
	HR(device->CreateRenderTargetView(texture, 0, &ssaoMapRTV1));
	HR(device->CreateShaderResourceView(texture, 0, &ssaoMapSRV1));
	HR(device->CreateUnorderedAccessView(texture, 0, &ssaoMapUAV1));

	//Done with the texture, release
	ReleaseCOM(texture);
}

void SSAOData::CreateSSAOQuad(ID3D11Device* device)
{
	SSAOMapFX::SSAOMapVertex vert[4];

	vert[0].PosL = XMFLOAT3(-1.0f, -1.0f, 0.0f);
	vert[1].PosL = XMFLOAT3(-1.0f, 1.0f, 0.0f);
	vert[2].PosL = XMFLOAT3(1.0f, 1.0f, 0.0f);
	vert[3].PosL = XMFLOAT3(1.0f, -1.0f, 0.0f);

	vert[0].Tex = XMFLOAT2(0.0f, 1.0f);
	vert[1].Tex = XMFLOAT2(0.0f, 0.0f);
	vert[2].Tex = XMFLOAT2(1.0f, 0.0f);
	vert[3].Tex  =XMFLOAT2(1.0f, 1.0f);

	vert[0].FarPlaneIndex = 0.0f;
	vert[1].FarPlaneIndex = 1.0f;
	vert[2].FarPlaneIndex = 2.0f;
	vert[3].FarPlaneIndex = 3.0f;

	D3D11_BUFFER_DESC bd;
	bd.BindFlags = D3D11_BIND_VERTEX_BUFFER;
	bd.ByteWidth = sizeof(SSAOMapFX::SSAOMapVertex) * 4;
	bd.CPUAccessFlags = 0;
	bd.MiscFlags = 0;
	bd.StructureByteStride = 0;
	bd.Usage = D3D11_USAGE_IMMUTABLE;

	D3D11_SUBRESOURCE_DATA sd;
	sd.pSysMem = vert;

	//Create vb
	HR(device->CreateBuffer(&bd, &sd, &ssaoMapGenVB));

	//Create index buffer.
	UINT ind[6] = 
	{
		0, 1, 2,
		0, 2, 3
	};

	bd.BindFlags = D3D11_BIND_INDEX_BUFFER;
	bd.ByteWidth = sizeof(UINT) * 6;
	
	sd.pSysMem = ind;

	HR(device->CreateBuffer(&bd, &sd, &ssaoMapGenIB));
}

void SSAOData::CreateOffsetVectors()
{
	// Start with 14 uniformly distributed vectors.  We choose the 8 corners of the cube
	// and the 6 center points along each cube face.  We always alternate the points on 
	// opposites sides of the cubes.  This way we still get the vectors spread out even
	// if we choose to use less than 14 samples.
	
	// 8 cube corners
	offsetVectors[0] = XMFLOAT4(+1.0f, +1.0f, +1.0f, 0.0f);
	offsetVectors[1] = XMFLOAT4(-1.0f, -1.0f, -1.0f, 0.0f);

	offsetVectors[2] = XMFLOAT4(-1.0f, +1.0f, +1.0f, 0.0f);
	offsetVectors[3] = XMFLOAT4(+1.0f, -1.0f, -1.0f, 0.0f);

	offsetVectors[4] = XMFLOAT4(+1.0f, +1.0f, -1.0f, 0.0f);
	offsetVectors[5] = XMFLOAT4(-1.0f, -1.0f, +1.0f, 0.0f);

	offsetVectors[6] = XMFLOAT4(-1.0f, +1.0f, -1.0f, 0.0f);
	offsetVectors[7] = XMFLOAT4(+1.0f, -1.0f, +1.0f, 0.0f);

	// 6 centers of cube faces
	offsetVectors[8] = XMFLOAT4(-1.0f, 0.0f, 0.0f, 0.0f);
	offsetVectors[9] = XMFLOAT4(+1.0f, 0.0f, 0.0f, 0.0f);

	offsetVectors[10] = XMFLOAT4(0.0f, -1.0f, 0.0f, 0.0f);
	offsetVectors[11] = XMFLOAT4(0.0f, +1.0f, 0.0f, 0.0f);

	offsetVectors[12] = XMFLOAT4(0.0f, 0.0f, -1.0f, 0.0f);
	offsetVectors[13] = XMFLOAT4(0.0f, 0.0f, +1.0f, 0.0f);

    for(int i = 0; i < 14; ++i)
	{
		// Create random lengths in [0.25, 1.0].
		float lenght = MathHelper::RandF(0.25f, 1.0f);
		
		XMVECTOR v = lenght * XMVector4Normalize(XMLoadFloat4(&offsetVectors[i]));
		
		XMStoreFloat4(&offsetVectors[i], v);
	}
}

void SSAOData::CreateRandomVectorTexture(ID3D11Device* device)
{
	ID3D11Texture2D* texture;

	D3D11_TEXTURE2D_DESC tDesc;
	tDesc.ArraySize = 1;
	tDesc.BindFlags = D3D11_BIND_SHADER_RESOURCE;
	tDesc.CPUAccessFlags = 0;
	tDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
	tDesc.Height = 256;
	tDesc.Width = 256;
	tDesc.MipLevels = 1;
	tDesc.MiscFlags = 0;
	tDesc.SampleDesc.Count = 1;
	tDesc.SampleDesc.Quality = 0;
	tDesc.Usage = D3D11_USAGE_IMMUTABLE;

	XMCOLOR initData[256*256];
	for (unsigned i = 0; i < 256; i++)
	{
		for (unsigned j = 0; j < 256; j++)
		{
			XMFLOAT3 vector(MathHelper::RandF(), MathHelper::RandF(), MathHelper::RandF());
			initData[i*256+j] = XMCOLOR(vector.x, vector.y, vector.z, 0.0f);
		}
	}

	//Create subresource data - ie, our initilisation data.
	D3D11_SUBRESOURCE_DATA data = {0};
	data.SysMemPitch = 256 * sizeof(XMCOLOR);
	data.pSysMem = initData;

	//Create texture.
	HR(device->CreateTexture2D(&tDesc, &data, &texture));
	//Create SRV
	HR(device->CreateShaderResourceView(texture, 0, &ssaoMapRandomVectorsSRV));

	//Release texture
	ReleaseCOM(texture);
}

void SSAOData::BuildFrustumCorners(float fovY, float farZ)
{
	float aspect = (float)rtWidth / (float)rtHeight;

	float halfHeight = farZ * tanf( 0.5f*fovY );
	float halfWidth  = aspect * halfHeight;

	frustumCorners[0] = XMFLOAT4(-halfWidth, -halfHeight, farZ, 0.0f);
	frustumCorners[1] = XMFLOAT4(-halfWidth, +halfHeight, farZ, 0.0f);
	frustumCorners[2] = XMFLOAT4(+halfWidth, +halfHeight, farZ, 0.0f);
	frustumCorners[3] = XMFLOAT4(+halfWidth, -halfHeight, farZ, 0.0f);
}